"""This module provides functionalities to work with image's shape.
Most of them are not implemented here yet but we can port and polish existing code here.
- change_band_axis
- compute_nonzero_bbox
- compute_nonzero_footprint
- crop
- flip
- rotate
- iter_image_by_block
- iter_image_by_block_with_buffer
"""
import enum
from typing import Tuple, Union, NamedTuple, Iterator, List, Optional, Literal

import numpy as np
import numba as nb
from shapely.geometry import Point, MultiPoint


class BandOrderType(str, enum.Enum):
    GDAL = "gdal"  # band at first axis
    IMAGE = "image"  # band at last axis


class BBoxInfo(NamedTuple):
    """
    Bounding Box Information in image row/col coordinate system

    Attributes
    ----------
    start_row : int
        Start row of the BBox.
    end_row : int
        End row of the BBox, non-inclusive.
    start_ind : int
        Start column of the BBox.
    end_ind : int
        End column of the BBox, non-inclusive.
    """

    start_row: int
    end_row: int
    start_col: int
    end_col: int


def change_band_axis(
    nda_img: np.ndarray,
    src_band_order: Union[str, BandOrderType] = BandOrderType.GDAL,
) -> np.ndarray:
    """Change the band axis location of a 3-dimensional image,
       from GDAL style (band at the first axis) to Image style (band at the last axis), or vice-versa.

    Parameters
    ----------
    nda_img : np.ndarray or ndarray-like
        A 3-dimension image in np.ndarray or ndarray like form
    src_band_order : Union[str, BandOrderType], optional
        Band order for the `nda_img`, can be one of "gdal", "image",
        BandOrderType.GDAL, or BandOrderType.IMAGE.
        By default BandOrderType.GDAL

    Returns
    -------
    np.ndarray or ndarray-like
        The 3-dimensional image in np.ndarray with band axis changed.
        If the given `nda_img` is a np.ndarray then a view of the original np.ndarray is returned,
        otherwise a new np.ndarray is returned.

    Raises
    ------
    ValueError
        Raised if the given `src_band_order` is not a known valid input.
    """
    if src_band_order == BandOrderType.GDAL:
        return np.swapaxes(nda_img, 0, -1)
    elif src_band_order == BandOrderType.IMAGE:
        return np.swapaxes(nda_img, -1, 0)
    else:
        raise ValueError(f"Unknown src_band_order {src_band_order}")


def compute_nonzero_bbox(nda_alpha: np.ndarray) -> Optional[BBoxInfo]:
    """Compute the nonzero area bounding box from a given 1-dimensional alpha ndarray

    Parameters
    ----------
    nda_alpha : np.ndarray or ndarray-like
        The input 1-dimensional alpha ndarray.
        Expect a bool ndarray, or an ndarray that implicitly uses 0 as the nodata mask value.

    Returns
    -------
    Optional[BBoxInfo]
        The detected nonzero bounding box. See `BBoxInfo` for the fields.
        or None if there is no nonzero zone within the given alpha ndarray.
        Notice the `end_row` and the `end_col` are non-inclusive.
    """
    rows, cols = np.nonzero(nda_alpha)
    if len(rows) == 0 or len(cols) == 0:
        return None

    row_min, row_max = np.min(rows), np.max(rows)
    col_min, col_max = np.min(cols), np.max(cols)
    return BBoxInfo(
        start_row=row_min,
        end_row=row_max + 1,
        start_col=col_min,
        end_col=col_max + 1,
    )


@nb.njit
def _get_possible_bounding_points(nda_alpha):
    """get all possible bounding points for a given 2D alpha array
    - assuming 0 indicates nodata and 1 indicates data
    - return a list of coordinates, each coordinate a tuple of (column_index, row_index) in image pixel space
    """
    res = []
    nrows, ncols = nda_alpha.shape

    # upper -> down, left side
    for i in range(nrows):
        for j in range(ncols):
            if nda_alpha[i, j] > 0:
                res.append((j, i))
                break

    # down -> upper, right side
    for i in range(nrows - 1, 0 - 1, -1):
        for j in range(ncols - 1, 0 - 1, -1):
            if nda_alpha[i, j] > 0:
                res.append((j, i))
                break
    return res


def _thin_bounding_points(all_bounding_points):
    """thinning a given list of bounding points by only keeping outer-most points"""
    nda_cols = np.array([item[0] for item in all_bounding_points])
    nda_rows = np.array([item[1] for item in all_bounding_points])

    bounding_points = set()

    # upper
    upper_row = nda_rows.min()
    for v in np.nonzero(nda_rows == upper_row)[0]:
        bounding_points.add((nda_cols[v], upper_row))

    # lower
    lower_row = nda_rows.max()
    for v in np.nonzero(nda_rows == lower_row)[0]:
        bounding_points.add((nda_cols[v], lower_row))

    # left
    left_col = nda_cols.min()
    for v in np.nonzero(nda_cols == left_col)[0]:
        bounding_points.add((left_col, nda_rows[v]))

    # right
    right_col = nda_cols.max()
    for v in np.nonzero(nda_cols == right_col)[0]:
        bounding_points.add((right_col, nda_rows[v]))

    return list(bounding_points)


def _reorg_bounding_points(bounding_points):
    """re-organize given bounding points in order and only keep points that form a convex hull"""
    tmp_max_row = max(item[1] for item in bounding_points)
    all_points = [Point(col, tmp_max_row - row) for col, row in bounding_points]
    all_points = MultiPoint(all_points)

    res = []
    for (col, row) in all_points.convex_hull.boundary.coords[:-1]:
        res.append((int(col), int(tmp_max_row - row)))
    return res


def compute_nonzero_convex_footprint(
    nda_alpha: np.ndarray, thin: bool = False
) -> List[Tuple[int, int]]:
    """Compute the nonzero area bounding box from a given 1-dimensional alpha ndarray

    Parameters
    ----------
    nda_alpha : np.ndarray or ndarray-like
        The input 1-dimensional alpha ndarray.
        Expect a bool ndarray, or an ndarray that implicitly uses 0 as the nodata mask value.
    thin: bool, optional
        Whether to thin the footprint list of coordinates
        By default False not to thin. See `Note` for more details.

    Returns
    -------
    List[Tuple[int, int]]
        The detected footprint as a list of coordinates,
        while each coordinate is a pair of two integers for the indices of column and row.
        or an empty list if there is no valid zone within the given alpha ndarray.

    Note
    ----
    This function assumes the valid portion with value 1 in the input `nda_alpha`
    is with a convex or near convex shape for its footprint.
    If the user has good confidence the footprint is convex and simple with only 4 corner points,
    the `thin` switch could be turned on to drastically speed up and simplify the result.
    """
    all_bounding_points = _get_possible_bounding_points(nda_alpha)
    if len(all_bounding_points) == 0:
        return []
    if thin:
        bounding_points = _thin_bounding_points(all_bounding_points)
    else:
        bounding_points = all_bounding_points
    pts = _reorg_bounding_points(bounding_points)
    return pts


def iter_image_by_block(
    nrows: int,
    ncols: int,
    block_nrows: int,
    block_ncols: int,
    boundary_treatment: str = "shrink",
) -> Iterator[BBoxInfo]:
    """Iterate through an image's coordinate space by block.

    Parameters
    ----------
    nrows : int
        Number of rows of the given image's coordinate space.
    ncols : int
        Number of columns of the given image's coordinate space.
    block_nrows : int
        Number of rows of the block's coordinate space.
    block_ncols : int
        Number of columns of the block's coordinate space.
    boundary_treatment : str, optional
        How to treat the boundary condition when the moving block meets the end of rows/columns.
        Either "shrink" to shrink the block so to be within the valid zone but with smaller size then specified,
        or "shift" to shift the block to the valid zone and keep the size to be the same as specified.
        By default "shrink".

    Yields
    -------
    Iterator[BBoxInfo]
        Bounding Box Info for the currently visited block. See `BBoxInfo` for the fields.
        Notice the `end_row` and the `end_col` are non-inclusive.

    Raises
    ------
    ValueError
        Raised if the given `boundary_treatment` is not one of "shrink" and "shift".
    """
    if boundary_treatment not in ("shrink", "shift"):
        raise ValueError("Unknown boundary_treatment {}".format(boundary_treatment))

    for ind_row in range(0, nrows, block_nrows):
        block_start_row = ind_row
        block_end_row = ind_row + block_nrows
        if block_end_row > nrows:
            if boundary_treatment == "shrink":
                block_end_row = nrows
            elif boundary_treatment == "shift":
                block_end_row = nrows
                block_start_row = max(nrows - block_nrows, 0)

        for ind_col in range(0, ncols, block_ncols):
            block_start_col = ind_col
            block_end_col = ind_col + block_ncols
            if block_end_col > ncols:
                if boundary_treatment == "shrink":
                    block_end_col = ncols
                elif boundary_treatment == "shift":
                    block_end_col = ncols
                    block_start_col = max(ncols - block_ncols, 0)
            yield BBoxInfo(
                block_start_row, block_end_row, block_start_col, block_end_col
            )


def iter_image_by_block_with_buffer(
    nrows: int,
    ncols: int,
    block_nrows: int,
    block_ncols: int,
    buffer_nrows: int,
    buffer_ncols: int,
    boundary_treatment: str = "shrink",
) -> Iterator[Tuple[BBoxInfo, BBoxInfo]]:
    """Iterate through an image's coordinate space by block with buffer.

    Parameters
    ----------
    nrows : int
        Number of rows of the given image's coordinate space.
    ncols : int
        Number of columns of the given image's coordinate space.
    block_nrows : int
        Number of rows of the block's coordinate space.
    block_ncols : int
        Number of columns of the block's coordinate space.
    buffer_nrows : int
        Buffer size for block's rows.
    buffer_ncols : int
        Buffer size for block's columns.
    boundary_treatment : str, optional
        How to treat the boundary condition when the moving block meets the end of rows/columns.
        Either "shrink" to shrink the block so to be within the valid zone but with smaller size then specified,
        or "shift" to shift the block to the valid zone and keep the size to be the same as specified.
        By default "shrink".

    Yields
    -------
    Iterator[Tuple[BBoxInfo, BBoxInfo]]
        Two Bounding Box Info for the currently visited block and buffered block.
        See `BBoxInfo` for the fields.
        Notice the `end_row` and the `end_col` are non-inclusive.

    Raises
    ------
    ValueError
        Raised if the given `boundary_treatment` is not one of "shrink" and "shift".
    """
    if boundary_treatment not in ("shrink", "shift"):
        raise ValueError("Unknown boundary_treatment {}".format(boundary_treatment))

    from common.img_utils.img_geom import _check_1D_boundary

    for valid_zone in iter_image_by_block(
        nrows, ncols, block_nrows, block_ncols, boundary_treatment=boundary_treatment
    ):
        # calculate the buffered zone based on the given zone and the buffer_nrows, buffer_ncols
        start_row, end_row = _check_1D_boundary(
            valid_zone.start_row - buffer_nrows,
            valid_zone.end_row + buffer_nrows,
            nrows,
            boundary_treatment,
        )
        start_col, end_col = _check_1D_boundary(
            valid_zone.start_col - buffer_ncols,
            valid_zone.end_col + buffer_ncols,
            ncols,
            boundary_treatment,
        )
        yield (valid_zone, BBoxInfo(start_row, end_row, start_col, end_col))


def valid_within_buffered(valid_block: BBoxInfo, buffered_block: BBoxInfo) -> BBoxInfo:
    """Compute the block info for the valid block within a given potentially larger buffered block

    Parameters
    ----------
    valid_block : BBoxInfo
        The small block representing the valid zone.
    buffered_block : BBoxInfo
        The potentially larger block representing the buffered zone.

    Returns
    -------
    BBoxInfo
        The block info representing the valid info within the buffered block.
    """

    start_row_tmp = valid_block.start_row - buffered_block.start_row
    end_row_tmp = start_row_tmp + (valid_block.end_row - valid_block.start_row)
    start_col_tmp = valid_block.start_col - buffered_block.start_col
    end_col_tmp = start_col_tmp + (valid_block.end_col - valid_block.start_col)

    return BBoxInfo(start_row_tmp, end_row_tmp, start_col_tmp, end_col_tmp)


@nb.jit(nopython=True, cache=True)
def patch_overlap_smoothing_for_boundary_shrink(
    nda_tile: np.array, nda_patch: np.array, buffered_block: BBoxInfo, buffer: int
) -> np.array:
    """A custom function to handle patch smoothing when iterating a large image with boundary `shrink`.
    NOTE, this function is NOT working for the boundary `shift`.

    Parameters
    ----------
    nda_tile : np.array
        The multi-channel image array in shape [C, H, W]., e.g. 10 x 10,980 x 10,980
    nda_patch : np.array
        The multi-channel patch array in shape [C, H, W], e.g., H = 1,000
    buffered_block : BBoxInfo
        The bounding box including buffer region for mapping nda_patch to nda_tile.
    buffer : int
        The buffer distance of smoothing.

    Returns
    -------
    np.array
        The patching result with weighted sum on overlapped region.
    """

    tile_width = nda_tile.shape[2]
    tile_height = nda_tile.shape[1]
    buffer_2x = buffer * 2

    for row in range(buffered_block.start_row, buffered_block.end_row, 1):
        # Check the row boundary
        delta_row = min(
            row - buffered_block.start_row, buffered_block.end_row - 1 - row
        )
        row_ratio = 1.0
        if delta_row < buffer_2x:
            # Disable smoothing on TILE boundary
            row_ratio = (
                delta_row / buffer_2x
                if (row > buffer_2x and row < (tile_height - 1 - buffer_2x))
                else 1.0
            )

        for col in range(buffered_block.start_col, buffered_block.end_col, 1):
            delta_col = min(
                col - buffered_block.start_col, buffered_block.end_col - 1 - col
            )
            col_ratio = 1.0
            # Only process pixels on PATCH boundary
            if delta_col < buffer_2x:
                # Disable smoothing on TILE boundary
                col_ratio = (
                    delta_col / buffer_2x
                    if (col > buffer_2x and col < (tile_width - 1 - buffer_2x))
                    else 1.0
                )

            # Only process pixels on PATCH boundary
            if col_ratio != 1.0 or row_ratio != 1.0:
                # Weighted Sum of pixel values on overlap regions
                nda_pixel = np.copy(
                    nda_patch[
                        :,
                        row - buffered_block.start_row,
                        col - buffered_block.start_col,
                    ]
                )
                nda_pixel = nda_pixel.astype(np.float32) * row_ratio * col_ratio
                nda_tile[:, row, col] += nda_pixel.astype(np.uint16)
            else:
                nda_tile[:, row, col] = nda_patch[
                    :, row - buffered_block.start_row, col - buffered_block.start_col
                ]

    return nda_tile


def block_mosaic(
    nda_lst: List[np.array],
    blocks_lst: List[BBoxInfo],
    image_size: Union[int, Tuple[int, int]],
    block_size: Union[int, Tuple[int, int]],
    buffer_size: Union[int, Tuple[int, int]],
    mosaic_mode: Literal["average", "gaussian", "idw", "bilinear"] = "average",
    mu: Optional[Union[float, int]] = None,
    sigma: Optional[Union[float, int]] = None,
) -> Tuple[np.array, np.array]:
    """A custom function to mosaic overlapping bloaks to a larger image.

    Parameters
    ----------
    nda_lst: List[np.array]
        the list of arrays that represent the multi-channel blocks (aka patches) in shape [C, H, W]
        the blocks in the list can be generated from RasterDatasetPerBlockTraverser, which returns a generator of
        blocks
    block_lst: List[BBoxInfo]
        the BBoxInfo corresponds to each block in the nda_lst
    image_size: Union[int, List[int, int]]
        the size of the raw image before being segmented into smaller blocks (aka patches), e.g., 2048 if the image
        has the same number of rows and columns, or [2048, 1024] if the image has different numbers of rows and
        columns
    block_size: Union[int, List[int, int]]
        the size of the blocks (aka patches) before to segment the raw image, e.g., 512 if the block has the same
        number of rows and columns, or [512, 256] if the block has different numbers of rows and columns
    buffer_size: Union[int, List[int, int]]
        the size of the buffer zone to add to each side of the block, e.g., 256 if to add 256 to top, bottom, left,
        right of the block, or [256, 128] if to add 256 to left and right but 128 to top and bottom
    mosaic_mode: Literal["average", "gaussian", "idw", "bilinear"] = "average"
        how to mosiac the overlaping region between patches, currently only support one of "average", "gaussian",
        "bilinear" and "idw"
        - "average": take the average of overlapping regions
        - "gaussian": the weight for overlapping regions follows a gaussian distribution
        - "idw": stands for inverse distance weight, the weight for overlapping regions depends on the distance
          to the block edge, the closer to the edge, the lower the weight
        - "bilinear": will do the same thing as "idw", but is more commonly used in computer vision
        by default "average"
    mu: numeric, optional
        the mean of the gaussian distribution
    sigma: numeric, optional
        the standard deviation of the gaussian distribution

    Returns
    -------
    weight: np.array
        the weight matrix to count the frequency of overlapping
    weighted_mosaic: np.array
        the merged tile taking weights into consideration
    """
    nrows, ncols = (
        image_size if isinstance(image_size, list) else image_size,
        image_size,
    )
    nrows_blk, ncols_blk = (
        block_size if isinstance(block_size, list) else block_size,
        block_size,
    )
    nrows_bfr, ncols_bfr = (
        buffer_size if isinstance(buffer_size, list) else buffer_size,
        buffer_size,
    )
    num_channels = nda_lst[0].shape[0]
    weight = np.zeros((nrows, ncols), dtype=np.float32)
    mosaic = np.zeros((num_channels, nrows, ncols), dtype=np.float32)
    original = np.zeros((num_channels, nrows, ncols), dtype=np.float32)
    img_center = (nrows_blk + 2 * nrows_bfr) // 2, (ncols_blk + 2 * ncols_bfr) // 2
    if mosaic_mode == "gaussian":
        x, y = np.meshgrid(
            range(ncols_blk + 2 * ncols_bfr), range(nrows_blk + 2 * nrows_bfr)
        )
        dst = np.sqrt((x - img_center[1]) ** 2 + (y - img_center[0]) ** 2)
        normal = 1 / (2.0 * np.pi * sigma ** 2)
        weight_patch_general = np.exp(-((dst - mu) ** 2 / (2.0 * sigma ** 2))) * normal
    elif mosaic_mode == "average":
        weight_patch_general = np.ones(
            [nrows_blk + 2 * nrows_bfr, ncols_blk + 2 * ncols_bfr]
        )
    elif mosaic_mode == "idw" or "bilinear":
        x, y = np.meshgrid(
            range(ncols_blk + 2 * ncols_bfr), range(nrows_blk + 2 * nrows_bfr)
        )
        x = np.where(x > (ncols_blk / 2 + ncols_bfr), ncols_blk + 2 * ncols_bfr - x, x)
        y = np.where(y > (nrows_blk / 2 + nrows_bfr), nrows_blk + 2 * nrows_bfr - y, y)
        # plus 0.1 is to aviod 0 at the margin
        weight_patch_general = (x + 0.1) * (y + 0.1)
    else:
        raise NotImplementedError(f"the mosaic mode {mosaic_mode} is not supported yet")
    for idx, img in enumerate(nda_lst):
        _, buffer_zone = blocks_lst[idx]
        radius_patch_row = (buffer_zone.end_row - buffer_zone.start_row) // 2
        radius_patch_col = (buffer_zone.end_col - buffer_zone.start_col) // 2
        weight_patch = weight_patch_general[
            img_center[0] - radius_patch_row : img_center[0] + radius_patch_row,
            img_center[1] - radius_patch_col : img_center[1] + radius_patch_col,
        ]
        mosaic[
            :,
            buffer_zone.start_row : buffer_zone.end_row,
            buffer_zone.start_col : buffer_zone.end_col,
        ] += (
            img * weight_patch[np.newaxis, :, :]
        )
        weight[
            buffer_zone.start_row : buffer_zone.end_row,
            buffer_zone.start_col : buffer_zone.end_col,
        ] += weight_patch
        original[
            :,
            buffer_zone.start_row : buffer_zone.end_row,
            buffer_zone.start_col : buffer_zone.end_col,
        ] = img
    return weight, (mosaic / weight[np.newaxis, :, :]).astype(np.uint16), original


# if __name__ == "__main__":
#     import numpy as np
#     from catalog.models.sentinel2_l3a import Sentinel2L3AScene
#     from catalog.raster.raster import retrieve_by_item
#     from common.geoimage.raster_dataset import RasterDatasetPerBlockTraverser
#     l3a_items = Sentinel2L3AScene.query_many_items(
#         temporal=["20220701", "20220709"],
#         mgrs_tile="49QCB"
#     )
#     item = l3a_items[0]
#     rds = retrieve_by_item(
#         item=item,
#         asset_names=["B04", "B03", "B02"]
#     )
#     patches = list(RasterDatasetPerBlockTraverser(
#         source_file_path=rds,
#         block_x_size=512,
#         block_y_size=512,
#         buffer_x_size=32,
#         buffer_y_size=32,
#         boundary_treatment="shift").traverse())
#     nda_lst_shift = [p[0].data for p in patches]
#     blocks_lst_shift = list(
#         iter_image_by_block_with_buffer(
#             nrows=rds.data.shape[1],
#             ncols=rds.data.shape[2],
#             block_nrows=512,
#             block_ncols=512,
#             buffer_nrows=32,
#             buffer_ncols=32,
#             boundary_treatment="shift",
#         )
#     )
#     weight_shift, mosaic_shift = block_mosaic(nda_lst_shift, blocks_lst_shift, rds.data.shape[1], 512, 32, "average")
